#include "MST.h"

#include <unordered_set>
#include <queue>
#include <cmath>
#include <algorithm>
#include <limits>
#include <memory>


void SteinerTreeSolver::computeSteinerTrees(
    std::unordered_map<std::string, std::vector<PinPolygon>>& nets,
    std::unordered_map<std::string, Via*>& net_vias,
    SteinerTreeMap& forest_roots,
    std::vector<std::vector<Line>>& flyLines) {

    flyLines.clear();
    forest_roots.clear();

    // 逐个计算每个net的斯坦纳树
    for (const auto& net_pair : nets) {
        const auto& net_name = net_pair.first;
        const auto& pins = net_pair.second;

        if (pins.size() <= 1) continue;

        // 获取对应的Via信息
        Via* via_template = net_vias[net_name];

        // 转换为指针集合（指向原始数据）
        std::vector<PinPolygon*> original_pin_ptrs;
        for (const auto& pin : pins) {
            original_pin_ptrs.emplace_back(const_cast<PinPolygon*>(&pin));
        }

        // 计算最小生成树（创建独立节点）
        SteinerTreePtr mst_root = computeMST(original_pin_ptrs);

        // 优化为斯坦纳树
        //mst_root = optimizeSteinerTree(mst_root, net_name, via_template);

        forest_roots[net_name] = mst_root;
    }

    // 遍历所有树，生成flyLines
    for (const auto& net_tree_pair : forest_roots) {
        SteinerTreePtr root = net_tree_pair.second;
        std::vector<Line> net_lines;
        generateFlyLinesFromTree(root, net_lines);
        flyLines.emplace_back(net_lines);
    }
}

// 计算最小生成树(Prim算法) - 返回独立创建的节点树
SteinerTreeSolver::SteinerTreePtr SteinerTreeSolver::computeMST(
    const std::vector<PinPolygon*>& original_pins) {

    if (original_pins.empty()) return nullptr;

    int n = original_pins.size();
    std::vector<bool> visited(n, false);
    std::vector<double> min_cost(n, std::numeric_limits<double>::max());
    std::vector<int> parent(n, -1);

    // 创建独立的节点副本（使用智能指针）
    std::vector<SteinerTreePtr> nodes;
    for (auto original_pin : original_pins) {
        PinPolygon* pin_copy = new PinPolygon(*original_pin);
        nodes.emplace_back(std::make_shared<SteinerNode>(pin_copy, true));
    }

    min_cost[0] = 0;

    for (int i = 0; i < n; ++i) {
        // 找到未访问的最小代价节点
        int u = -1;
        double min_val = std::numeric_limits<double>::max();
        for (int j = 0; j < n; ++j) {
            if (!visited[j] && min_cost[j] < min_val) {
                min_val = min_cost[j];
                u = j;
            }
        }

        if (u == -1) break;
        visited[u] = true;

        // 构建父子关系
        if (parent[u] != -1) 
            nodes[parent[u]]->addChild(nodes[u]);

        // 更新邻接节点
        for (int v = 0; v < n; ++v) {
            if (!visited[v]) {
                double dist = nodes[u]->position.distanceTo(nodes[v]->position);
                if (dist < min_cost[v]) {
                    min_cost[v] = dist;
                    parent[v] = u;
                }
            }
        }
    }
    return nodes[0]; // 返回根节点
}

// 从树生成连线
void SteinerTreeSolver::generateFlyLinesFromTree(SteinerTreePtr root, std::vector<Line>& flyLines) {
    if (!root) return;

    std::queue<SteinerTreePtr> q;
    std::unordered_set<SteinerTreePtr> visited;

    q.push(root);
    visited.insert(root);

    while (!q.empty()) {
        SteinerTreePtr current = q.front();
        q.pop();

        // 添加与所有子节点的连线
        for (auto child : current->children) {
            if (visited.find(child) == visited.end()) {
                flyLines.emplace_back(current->position, child->position);
                q.push(child);
                visited.insert(child);
            }
        }

        // 添加与父节点的连线（避免重复）
        if (auto parent_ptr = current->parent.lock()) {
            if (visited.find(parent_ptr) == visited.end()) {
                flyLines.emplace_back(current->position, parent_ptr->position);
            }
        }
    }
}

/*

// 插入斯坦纳点，优化斯坦纳树
SteinerTreeSolver::SteinerTreePtr SteinerTreeSolver::optimizeSteinerTree(
    SteinerTreePtr root,
    const std::string& net_name,
    const Via& via_template) {

    if (!root) return nullptr;

    std::vector<SteinerTreePtr> all_nodes;
    collectAllNodes(root, all_nodes);

    // 如果节点数足够，尝试添加斯坦纳点
    if (all_nodes.size() >= 3) {
        // 1.计算所有点的中心
        Point center;
        getMidPoint(all_nodes, center);

        // 2.创建斯坦纳点
        PinPolygon* steiner_pin = createSteinerPoint(center, all_nodes, via_template);
        m_extraPolygons[net_name].emplace_back(steiner_pin);

        SteinerTreePtr steiner_node = std::make_shared<SteinerNode>(steiner_pin, false);

        // 找到距离斯坦纳点最近的3个节点
        std::vector<std::pair<double, SteinerTreePtr>> distances;
        for (auto node : all_nodes) {
            double dist = center.distanceTo(node->position);
            distances.emplace_back({ dist, node });
        }

        std::sort(distances.begin(), distances.end());

        // 重新构建连接关系
        for (int i = 0; i < std::min(3, (int)distances.size()); ++i) {
            SteinerTreePtr nearest_node = distances[i].second;

            // 如果最近节点有父节点，先断开连接
            if (auto parent_ptr = nearest_node->parent.lock()) {
                nearest_node->disconnectFromParent();
            }

            // 连接到斯坦纳点
            steiner_node->addChild(nearest_node);
        }

        return steiner_node;
    }

    return root;
}

// 收集所有节点（BFS遍历）
void SteinerTreeSolver::collectAllNodes(SteinerTreePtr root, std::vector<SteinerTreePtr>& nodes) {
    if (!root) return;

    std::queue<SteinerTreePtr> q;
    std::unordered_set<SteinerTreePtr> visited;

    q.push(root);
    visited.insert(root);
    nodes.emplace_back(root);

    while (!q.empty()) {
        SteinerTreePtr current = q.front();
        q.pop();

        for (auto child : current->children) {
            if (visited.find(child) == visited.end()) {
                visited.insert(child);
                nodes.emplace_back(child);
                q.push(child);
            }
        }
    }
}

void SteinerTreeSolver::getMidPoint(const std::vector<SteinerTreePtr>& all_nodes, Point& center) {
    // 规则1：计算所有点的几何中心（默认规则）
    center.x = 0;
    center.y = 0;
    for (auto node : all_nodes) {
        center.x += node->position.x;
        center.y += node->position.y;
    }
    center.x /= all_nodes.size();
    center.y /= all_nodes.size();
}

// 创建斯坦纳点
PinPolygon* SteinerTreeSolver::createSteinerPoint(
    const Point& position,
    const std::vector<SteinerTreePtr>& neighbors,
    const Via& via_template) {

    // 收集邻居的所有layer
    std::unordered_set<int> layer_set;
    for (auto neighbor : neighbors) {
        for (int layer : neighbor->pin->m_layers) {
            layer_set.insert(layer);
        }
    }
    std::vector<int> layers(layer_set.begin(), layer_set.end());

    std::vector<Line> edges;
    PinPolygon* new_pin = new PinPolygon(const_cast<Point&>(position), layers, edges, via_template.m_radius);
    return new_pin;
}

*/